# -*- coding: UTF-8 -*-

from __future__ import absolute_import, division, print_function, unicode_literals

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])

# Import the backtrader platform
import backtrader as bt

import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import talib as talib
import numpy as np
import akshare as ak
import pandas as pd
import json


class TrendStrategy(bt.Strategy):
    params = (
        ("ema_short", 20),
        ("ema_long", 60),
        ("vol_ma", 5),
    )

    def __init__(self):
        # 计算技术指标
        self.ema20 = bt.indicators.EMA(period=self.p.ema_short)
        self.ema60 = bt.indicators.EMA(period=self.p.ema_long)
        self.macd = bt.indicators.MACD()

        # 量能均线
        self.vol_ma = bt.indicators.SMA(self.data.volume, period=self.p.vol_ma)

        # 动态支撑压力
        self.support, self.resistance = self.calc_support_resistance()

    def calc_support_resistance(self):
        return self.data.close[0], self.data.close[0]
        # high = self.data.high.get(size=20)
        # low = self.data.low.get(size=20)

        # print(
        #     "high = ",
        #     high,
        #     ", low = ",
        #     low,
        #     ", len(high) = ",
        #     len(high),
        #     ", len(low) = ",
        #     len(low),
        #     "self.data.close=",
        #     self.data.close[0],
        # )

        # if len(high) < 20 or len(low) < 20:
        #     return self.data.close[0], self.data.close[0]

        # pivot = (
        #     max(self.data.close.get(size=20)) + min(self.data.close.get(size=20))
        # ) / 2
        # return pivot * 0.9, pivot * 1.1
        # pivot = (max(high) + min(low)) / 2
        # return pivot * 0.9, pivot * 1.1

    def next(self):
        # 趋势判断
        trend_up = self.ema20 > self.ema60 and self.data.close[0] > self.resistance

        # 量能验证
        vol_cond = self.data.volume[0] > 1.5 * self.vol_ma[0]

        # 买卖信号
        if not self.position:
            if trend_up and vol_cond:
                self.buy(size=self.broker.getvalue() // self.data.close[0] * 0.2)
        else:
            if self.data.close < self.support:
                self.sell()


# 获取股票数据
def get_stock_data(code):
    df = ak.stock_zh_a_hist(symbol=code, adjust="hfq")
    df = df[["日期", "开盘", "最高", "最低", "收盘", "成交量"]]
    df.columns = ["date", "open", "high", "low", "close", "volume"]
    return df


if __name__ == "__main__":
    code = "002202"
    # df = get_stock_data(code)
    # sector_df = ak.stock_board_industry_hist_em()
    # data = bt.feeds.PandasData(dataname=df)

    cerebro = bt.Cerebro()

    # 获取行业板块数据
    tableName = "t_trade_data_" + code
    sql = f"select `date`, `open`, `high`, `low`, `close`, `volume` from {tableName}"

    data = mdf.MySQLDataFeed(sql)
    print("data = ", data)
    # Add the Data Feed to Cerebro
    cerebro.adddata(data, name=code)
    # 创建策略对象
    cerebro.adddata(data)
    cerebro.addstrategy(TrendStrategy)
    cerebro.broker.setcash(1000000)
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe")
    cerebro.run()
    cerebro.plot()
